﻿#include "Mp4Encoder.h"
#include <memory>

#pragma comment(lib,"libmp4v2.lib")

Mp4Encoder::Mp4Encoder()
	: _videoWidth(640),
	_videoHeight(480),
	_videoTimeScale(10000000),
	_mp4FHandle(nullptr),
	_audioTrackId(MP4_INVALID_TRACK_ID),
	_videoTrackId(MP4_INVALID_TRACK_ID)
{

}

bool Mp4Encoder::initMp4Encoder(std::string& config, int sampleNum, int sampleRate)
{
	if (_mp4FHandle)
	{
		return true;
	}

	if (!_filePathName.size())
	{
		return false;
	}

	//-------------------- file handle --------------------
	_mp4FHandle = MP4Create(_filePathName.c_str());
	if (_mp4FHandle == MP4_INVALID_FILE_HANDLE)
	{
		return false;
	}
	MP4SetTimeScale(_mp4FHandle, _videoTimeScale);

	//-------------------- audio track --------------------
	if (config.size())
	{
		_audioTrackId = MP4AddAudioTrack(_mp4FHandle, sampleRate, sampleNum, MP4_MPEG4_AUDIO_TYPE);
		if (_audioTrackId == MP4_INVALID_TRACK_ID)
		{
			return false;
		}

		MP4SetAudioProfileLevel(_mp4FHandle, 0x2);
		MP4SetTrackESConfiguration(_mp4FHandle, _audioTrackId, (uint8_t*)config.c_str(), config.size());
	}
	return true;
}

//------------------------------ h264码流分析 ------------------------------------------
//	h264 NALU:  0x00 00 00 01 | nalu_type(1字节)| nalu_data (N 字节) | 0x00 00 00 01 | ...
//				起始码(4字节)          类型             数据           下一个NALU起始码
//  nalu_type有以下种类，长度为1个字节，8位分别表示：【 1比特 禁止位，恒为0 | 2比特 重要性指示位 | 5比特 类型 】
//	0x67 (0 11 00111) SPS		非常重要			
//	0x68 (0 11 01000) PPS		非常重要			
//	0x65 (0 11 00101) IDR帧		非常重要，关键帧
//	0x61 (0 11 00001) I帧		重要，非IDR的关键帧
//	0x41 (0 10 00001) P帧		重要				 
//	0x01 (0 00 00001) B帧		不重要				
//	0x06 (0 00 00110) SEI		不重要				
// 这里对每个NAL单元的前4或5个字节进行解析，从而查找是什么类型的NAL单元
// 第一帧 SPS【0 0 0 1 0x67】 PPS【0 0 0 1 0x68】 SEI【0 0 1 0x6】 IDR【0 0 1 0x65】
// p帧      P【0 0 0 1 0x41】
// I帧    SPS【0 0 0 1 0x67】 PPS【0 0 0 1 0x68】 IDR【0 0 1 0x65】
// 除非编码器和解码器进行特定的语法协商，解码器一般不对SEI包进行解析，所以这里不处理SEI数据
//
// MP4WriteSample接收I、P帧 nalu,该nalu需要用4字节的数据大小头替换原有的起始头，并且数据大小为big-endian格式（即将高字节存储在起始地址）
//--------------------------------------------------------------------------------------

void Mp4Encoder::mp4VEncode(unsigned char* naluData, int naluSize, double takeTime)
{
	int index = -1;

	if (naluData[0] == 0 && naluData[1] == 0 && naluData[2] == 0 && naluData[3] == 1 && naluData[4] == 0x67)
	{
		index = 0x67;
	}
	if (index != 0x67 && _videoTrackId == MP4_INVALID_TRACK_ID)
	{
		return;
	}
	if (naluData[0] == 0 && naluData[1] == 0 && naluData[2] == 0 && naluData[3] == 1 && naluData[4] == 0x68)
	{
		index = 0x68;
	}
	if (naluData[0] == 0 && naluData[1] == 0 && naluData[2] == 1 && naluData[3] == 0x65)
	{
		index = 0x65;
	}
	if (naluData[0] == 0 && naluData[1] == 0 && naluData[2] == 0 && naluData[3] == 1 && naluData[4] == 0x41)
	{
		index = 0x41;
	}
	if (naluData[0] == 0 && naluData[1] == 0 && naluData[2] == 0 && naluData[3] == 1 && naluData[4] == 0x01)
	{
		index = 0x01;
	}

	switch (index)
	{
	case 0x67: // SPS 
		if (_videoTrackId == MP4_INVALID_TRACK_ID)
		{
			_videoTrackId = MP4AddH264VideoTrack(
				_mp4FHandle,
				_videoTimeScale,
				MP4_INVALID_DURATION,
				_videoWidth,
				_videoHeight,
				naluData[5], // sps[1] AVCProfileIndication  
				naluData[6], // sps[2] profile_compat  
				naluData[7], // sps[3] AVCLevelIndication  
				3);          // 4 bytes length before each NAL unit 

			if (_videoTrackId == MP4_INVALID_TRACK_ID)
			{
				return;
			}
			MP4SetVideoProfileLevel(_mp4FHandle, 0x7F);
		}

		MP4AddH264SequenceParameterSet(_mp4FHandle, _videoTrackId, naluData + 4, naluSize - 4);
		break;
	case 0x68: // PPS
		MP4AddH264PictureParameterSet(_mp4FHandle, _videoTrackId, naluData + 4, naluSize - 4);
		break;
	case 0x65: // I
	{
		unsigned char* data = new unsigned char[naluSize + 1];
		data[0] = (naluSize - 3) >> 24;
		data[1] = (naluSize - 3) >> 16;
		data[2] = (naluSize - 3) >> 8;
		data[3] = (naluSize - 3) & 0xff;
		memcpy(data + 4, naluData + 3, naluSize - 3);

		// 注意这里第五个参数，是指两个视频帧之间的tick数目，takeTime是时间间隔
		MP4WriteSample(_mp4FHandle, _videoTrackId, data, naluSize + 1, takeTime, 0, 1);

		delete[] data;
		break;
	}
	case 0x41: // P
	{
		naluData[0] = (naluSize - 4) >> 24;
		naluData[1] = (naluSize - 4) >> 16;
		naluData[2] = (naluSize - 4) >> 8;
		naluData[3] = (naluSize - 4) & 0xff;

		// 注意这里第五个参数，是指两个视频帧之间的tick数目，takeTime是时间间隔
		MP4WriteSample(_mp4FHandle, _videoTrackId, naluData, naluSize, takeTime, 0, 1);

		break;
	}
	case 0x01:
	{
		break;
	}
	default:
		break;
	}
}

void Mp4Encoder::mp4AEncode(unsigned char* aacData, int aacSize)
{
	if (_videoTrackId == MP4_INVALID_TRACK_ID)
	{
		return;
	}

	// 音频只要和采样时间相近的视频帧一同写入即可
	MP4WriteSample(_mp4FHandle, _audioTrackId, aacData, aacSize, MP4_INVALID_DURATION, 0, 1);
}

void Mp4Encoder::closeMp4Encoder()
{
	if (_mp4FHandle)
	{
		MP4Close(_mp4FHandle);
		_mp4FHandle = nullptr;
	}

	_audioTrackId = MP4_INVALID_TRACK_ID;
	_videoTrackId = MP4_INVALID_TRACK_ID;
}